# -*- coding: utf-8 -*-
"""
@author:XuMing(xuming624@qq.com)
@description:
"""

"""Natural Language Inference (NLI) Chinese Corpus.(nli_zh)"""

import os
import json
import datasets

_DESCRIPTION = """SimCLUE：3000000+中文语义理解与匹配数据集"""

GITHUB_HOME = "https://github.com/CLUEbenchmark/SimCLUE"

_CITATION = "https://github.com/CLUEbenchmark/SimCLUE"

_DATA_URL = "https://storage.googleapis.com/cluebenchmark/tasks/simclue_public.zip"


class SimCLUEConfig(datasets.BuilderConfig):

    def __init__(self, features, data_url, citation, url, label_classes=(0, 1), **kwargs):
        super().__init__(version=datasets.Version("1.0.0"), **kwargs)
        self.features = features
        self.label_classes = label_classes
        self.data_url = data_url
        self.citation = citation
        self.url = url


class SimCLUE(datasets.GeneratorBasedBuilder):
    """The Natural Language Inference Chinese(NLI_zh) Corpus."""

    part_file = {'train_rank': 'train_rank.json',
                 'train_pair': 'train_pair.json',
                 'corpus': 'corpus.txt',
                 'train_pair_postive': 'train_pair_postive.json',
                 'dev': 'dev.json',
                 'test_public': 'test_public.json'}

    part_split = {'train_rank': datasets.Split.TRAIN,
                 'train_pair': datasets.Split.TRAIN,
                 'corpus': datasets.Split.TRAIN,
                 'train_pair_postive': datasets.Split.TRAIN,
                 'dev': datasets.Split.VALIDATION,
                 'test_public': datasets.Split.TEST}

    BUILDER_CONFIGS = [
        SimCLUEConfig(
            name="train_rank",
            description=_DESCRIPTION,
            features=datasets.Features({"query": datasets.Value("string"),
                                          "title": datasets.Value("string"),
                                          "neg_title": datasets.Value("string")}),
            data_url=_DATA_URL,
            citation=_CITATION,
            url=GITHUB_HOME,
        ),
        SimCLUEConfig(
            name="train_pair",
            description=_DESCRIPTION,
            features=datasets.Features({"sentence1": datasets.Value("string"),
                                       "sentence2": datasets.Value("string"),
                                       "label": datasets.Value("int32")}),
            data_url=_DATA_URL,
            citation=_CITATION,
            url=GITHUB_HOME,
        ),
        SimCLUEConfig(
            name="corpus",
            description=_DESCRIPTION,
            features=datasets.Features({"sentence1": datasets.Value("string")}),
            data_url=_DATA_URL,
            citation=_CITATION,
            url=GITHUB_HOME,
        ),
        SimCLUEConfig(
            name="train_pair_postive",
            description=_DESCRIPTION,
            features=datasets.Features({"sentence1": datasets.Value("string"),
                                          "sentence2": datasets.Value("string"),
                                          "label": datasets.Value("int32")}),
            data_url=_DATA_URL,
            citation=_CITATION,
            url=GITHUB_HOME,
        ),
        SimCLUEConfig(
            name="dev",
            description=_DESCRIPTION,
            features=datasets.Features({"sentence1": datasets.Value("string"),
                                          "sentence2": datasets.Value("string"),
                                          "label": datasets.Value("int32")}),
            data_url=_DATA_URL,
            citation=_CITATION,
            url=GITHUB_HOME,
        ),
        SimCLUEConfig(
            name="test_public",
            description=_DESCRIPTION,
            features=datasets.Features({"sentence1": datasets.Value("string"),
                                          "sentence2": datasets.Value("string"),
                                          "label": datasets.Value("int32")}),
            data_url=_DATA_URL,
            citation=_CITATION,
            url=GITHUB_HOME,
        ),
    ]

    def _info(self):
        return datasets.DatasetInfo(
            description=self.config.description,
            features=self.config.features,
            homepage=self.config.url,
            citation=self.config.citation,
        )

    def _split_generators(self, dl_manager):
        dl_dir = dl_manager.download_and_extract(self.config.data_url)
        return [datasets.SplitGenerator(
                        name=self.part_split[self.config.name],
                        gen_kwargs={
                            "filepath": os.path.join(dl_dir, self.part_file[self.config.name]),
                        })]

    def _generate_examples(self, filepath):
        """This function returns the examples in the raw (text) form."""
        with open(filepath, 'r', encoding="utf-8") as f:
            for idx, row in enumerate(f):
                yield idx, json.loads(row)